package sample.methodremoveinstruct;

import org.objectweb.asm.*;
import util.FileUtils;

import java.io.IOException;

import static org.objectweb.asm.Opcodes.*;
import static util.StateMachineUtils.SEEN_NOTHING;

public class MethodRemoveGetFieldPutFieldDemo {
    public static void main(String[] args) throws IOException {
        Class<?> clazz = HelloWorld.class;
        byte[] bytes = invoke(clazz);
        //  output
        FileUtils.writeBytes(Type.getInternalName(clazz) + "RemoveGetFieldPutField.class", bytes);
        //  replace
        FileUtils.writeBytes(clazz.getResource("/").getPath() + Type.getInternalName(clazz) + ".class", bytes);
    }

    private static byte[] invoke(Class<?> clazz) throws IOException {
        ClassReader classReader = new ClassReader(clazz.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES);
        MethodRemoveGetFieldPutFieldVisitor visitor = new MethodRemoveGetFieldPutFieldVisitor(Opcodes.ASM9, classWriter);
        classReader.accept(visitor, ClassReader.EXPAND_FRAMES);
        return classWriter.toByteArray();

    }

    protected static class MethodRemoveGetFieldPutFieldVisitor extends ClassVisitor {
        private final int SEEN_ALOAD_0 = 1;
        private final int SEEN_ALOAD_0_ALOAD_0 = 2;
        private final int SEEN_ALOAD_0_ALOAD_0_GETFIELD = 3;

        public MethodRemoveGetFieldPutFieldVisitor(int api, ClassVisitor classVisitor) {
            super(api, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            if (mv != null && !"<init>".equals(name) && !"<clinit>".equals(name)) {
                boolean isAbstractMethod = (access & ACC_ABSTRACT) != 0;
                boolean isNativeMethod = (access & ACC_NATIVE) != 0;
                if (!isAbstractMethod && !isNativeMethod) {
                    mv = new MethodRemoveGetFieldPutFieldAdapter(api, mv);
                }
            }
            return mv;
        }

        private class MethodRemoveGetFieldPutFieldAdapter extends MethodPatternAdapter {

            private String fieldOwner;
            private String fieldName;
            private String fieldDesc;

            public MethodRemoveGetFieldPutFieldAdapter(int api, MethodVisitor mv) {
                super(api, mv);
            }

            @Override
            public void visitVarInsn(int opcode, int var) {
                if (state == SEEN_NOTHING && opcode == ALOAD && var == 0) {
                    state = SEEN_ALOAD_0;
                    return;
                }
                if (state == SEEN_ALOAD_0 && opcode == ALOAD && var == 0) {
                    state = SEEN_ALOAD_0_ALOAD_0;
                    return;
                }
                if (state == SEEN_ALOAD_0_ALOAD_0 && opcode == ALOAD && var == 0) {
                    super.visitVarInsn(opcode, var);
                    return;
                }
                super.visitVarInsn(opcode, var);
            }

            @Override
            public void visitFieldInsn(int opcode, String owner, String name, String descriptor) {
                if (state == SEEN_ALOAD_0_ALOAD_0 && opcode == GETFIELD) {
                    state = SEEN_ALOAD_0_ALOAD_0_GETFIELD;
                    this.fieldOwner = owner;
                    this.fieldName = name;
                    this.fieldDesc = descriptor;
                    return;
                }
                if (state == SEEN_ALOAD_0_ALOAD_0_GETFIELD && opcode == PUTFIELD && name.equals(fieldName)) {
                    state = SEEN_NOTHING;
                    return;
                }
                super.visitFieldInsn(opcode, owner, name, descriptor);
            }

            @Override
            protected void visitInsn() {
                switch (state) {
                    case SEEN_ALOAD_0:
                        mv.visitVarInsn(ALOAD, 0);
                        break;
                    case SEEN_ALOAD_0_ALOAD_0:
                        mv.visitVarInsn(ALOAD, 0);
                        mv.visitVarInsn(ALOAD, 0);
                        break;
                    case SEEN_ALOAD_0_ALOAD_0_GETFIELD:
                        mv.visitVarInsn(ALOAD, 0);
                        mv.visitVarInsn(ALOAD, 0);
                        mv.visitFieldInsn(GETFIELD, fieldOwner, fieldName, fieldDesc);
                        break;
                    default:
                        break;
                }
                state = SEEN_NOTHING;
            }
        }
    }
}
